library(R.matlab)
library(ggplot2)
library(ggforce)
library(tidyverse)
library(mvtnorm)
library(grid)
library(svglite)
library(showtext)
library(cowplot)
library(Cairo)

pth_reg <- "/System/Library/Fonts/Palatino.ttc"
font_add(family = "Palatino", regular = pth_reg)


t <- 0.134
epsilon <- 0.000094
w <-  1.0 


setwd("/Users/jbhoward/Documents/Personal Research/JAERE Extension/FINAL SIMULATION")
out_path <- file.path("figures", "output")


fpth <- file.path("simulation", "counterfactual_results.mat")
mat <- readMat(fpth)

all_results <- mat$endog.results

objects <- all_results[, , 1]

phi <- objects$phi
del <- objects$del
n_grid <- length(phi)


g_legend <- function(a.gplot){
  tmp <- ggplot_gtable(ggplot_build(a.gplot))
  leg <- which(sapply(tmp$grobs, function(x) x$name) == "guide-box")
  legend <- tmp$grobs[[leg]]
  return(legend)
}


############## GET GCHI DATA!!!! ##################
M = c(mat$m.phi, mat$m.del)                                        
V = matrix(c(mat$v.phi^2, mat$v.cor*mat$v.phi*mat$v.del, mat$v.cor*mat$v.phi*mat$v.del, mat$v.del^2), ncol = 2) 
         
g_chi_fun <- function(phi,del) dmvnorm(c(log(phi), log(del)), M, sigma =V, log = FALSE, checkSymmetry = TRUE)/(phi*del)
g_chi <- matrix(rep(0, n_grid*n_grid), ncol = n_grid)
for (n_phii in seq_len(n_grid)) {
    for (n_dell in seq_len(n_grid)) {
        g_chi[n_phii, n_dell] <- g_chi_fun(phi[n_phii], del[n_dell])
    }
}
g_chi <- g_chi/sum(g_chi)


# initialize our data frame
df <- data.frame("phi" = rep(phi, n_grid),
                "del" = rep(del, each = n_grid),
                "g_chi" = c(g_chi))

vars <- c("prob.coastal", "prob.inland",
    "prob.inland.domestic", "prob.inland.exporter",
    "prob.coastal.domestic", "prob.coastal.exporter",
    "opt.e.coastal", "opt.e.inland",
    "inland.X", "coastal.X", "d.inland", "d.coastal",
    "share.fa.inland", "share.fa.coastal",
    "dxratio.inland", "dxratio.coastal", "total.profit",
    "xdratio.inland", "xdratio.coastal", "p.x", "Phi.inland", "Phi.coastal",
    "fa.inland", "fa.coastal", "int.x.coastal", "int.x.inland",
    "share.d.inland", "share.d.coastal", "Kappa.inland", "Kappa.coastal")

for (var in vars) {
    df[var] <- c(objects[[var]])
}

# construct "intensity" manually
df$coastal.eint <- df$opt.e.coastal / df$coastal.X
df$inland.eint <- df$opt.e.inland / df$inland.X


loc_vars <- c("prob.inland", "prob.coastal")

df$loc_prob_max <- apply(df[, loc_vars], 1, max, na.rm=TRUE)
df$firm_loc <- colnames(df[, loc_vars])[apply(df[, loc_vars], 1, which.max)]

df$firm_loc <- "prob.coastal"
il_prob <- mean(df$prob.inland)
df$firm_loc[df$prob.inland >= il_prob] <- "prob.inland"

df$firm_loc[df$loc_prob_max == 0] <- "Exit"
# df$firm_loc[df$total.profit < 0.01] <- "Exit"

intl_vars <- c("prob.inland.domestic", "prob.inland.exporter")
df$intl_inland <- colnames(df[, intl_vars])[apply(df[, intl_vars], 1, which.max)]

intl_vars <- c("prob.coastal.domestic", "prob.coastal.exporter")
df$intl_coastal <- colnames(df[, intl_vars])[apply(df[, intl_vars], 1, which.max)]

cutoff <- mean(df$prob.coastal.exporter[df$firm_loc == "prob.coastal"])
df$intl_coastal <- "prob.coastal.exporter"
df$intl_coastal[df$prob.coastal.exporter < cutoff] <- "prob.coastal.domestic"

df$firmtype <- "Exit"
df$firmtype[df$firm_loc == "prob.inland"] <- df$intl_inland[df$firm_loc == "prob.inland"]
df$firmtype[df$firm_loc == "prob.coastal"] <- df$intl_coastal[df$firm_loc == "prob.coastal"]


df$firmtype[df$firmtype == "prob.inland.domestic"] <- "PD"
df$firmtype[df$firmtype == "prob.inland.exporter"] <- "PI"
df$firmtype[df$firmtype == "prob.coastal.domestic"] <- "CD"
df$firmtype[df$firmtype == "prob.coastal.exporter"] <- "CI"

df$m <- df$max

df$type = "Exit"
df$type[df$firmtype == "PD"] <- "Periphery-Domestic"
df$type[df$firmtype == "PI"] <- "Periphery-International"
df$type[df$firmtype == "CI"] <- "Core-International"
df$type[df$firmtype == "CD"] <- "Core-Domestic"


# get quantiles of phi and del
phis <- unique(df$phi)
phi_qt <- data.frame(phi = phis, phi_qt = seq_along(phis))
dels <- unique(df$del)
del_qt <- data.frame(del = dels, del_qt = seq_along(dels))
df <- merge(df, phi_qt, by = "phi")
df <- merge(df, del_qt, by = "del")

scale <- 100 / max(df$phi_qt, df$del_qt)

phi_pts <- c()
for (del in unique(df$del_qt)) {
    v1 <- min(df[df$del_qt == del & df$firmtype == "CI", "phi_qt"])
    v2 <- max(df[df$del_qt == del & df$firmtype == "CD", "phi_qt"])
    phi_pts <- c(phi_pts, mean(v1, v2))
}
del_pts_grid <- unique(df$del_qt)

# draw line between cd and ci 
del_pts <- c()
for (phi in unique(df$phi_qt)) {
    v1 <- min(df[df$phi_qt == phi & df$firmtype == "CI", "del_qt"])
    v2 <- max(df[df$phi_qt == phi & df$firmtype == "CD", "del_qt"])
    del_pts <- c(del_pts, mean(v1, v2))
}
phi_pts_grid <- unique(df$phi_qt)

lines_1 <- data.frame(del_qt = del_pts, phi_qt = phi_pts_grid)
lines_2 <- data.frame(del_qt = del_pts_grid, phi_qt = phi_pts)
lines_cd <- unique(union(lines_1, lines_2))
lines_cd <- lines_cd[order(lines_cd$del_qt, -lines_cd$phi_qt), ]
lines_cd$del_qt <- lines_cd$del_qt - 0.5
lines_cd$phi_qt <- lines_cd$phi_qt - 0.5

# draw line between pd and cd
phi_pts <- c()
for (del in unique(df$del_qt)) {
    v1 <- min(df[df$del_qt == del & df$firmtype == "CD", "phi_qt"])
    v2 <- max(df[df$del_qt == del & df$firmtype == "PD", "phi_qt"])
    phi_pts <- c(phi_pts, mean(v1, v2))
}

del_pts_grid <- unique(df$del_qt)

# draw line between pd and cd 
del_pts <- c()
for (phi in unique(df$phi_qt)) {
    v1 <- min(df[df$phi_qt == phi & df$firmtype == "CD", "del_qt"])
    v2 <- max(df[df$phi_qt == phi & df$firmtype == "PD", "del_qt"])
    del_pts <- c(del_pts, mean(v1, v2))
}

phi_pts_grid <- unique(df$phi_qt)
lines_1 <- data.frame(del_qt = del_pts, phi_qt = phi_pts_grid)
lines_2 <- data.frame(del_qt = del_pts_grid, phi_qt = phi_pts)
lines_ci <- unique(union(lines_1, lines_2))
lines_ci <- lines_ci[order(lines_ci$del_qt, -lines_ci$phi_qt), ]
lines_ci$del_qt <- lines_ci$del_qt - 0.5
lines_ci$phi_qt <- lines_ci$phi_qt - 0.5

max_phi <- max(df[df$firmtype == "PD", "phi_qt"])+1
max_del <- max(df[df$firmtype == "PD", "del_qt"])+1

lines_ci <- lines_ci[lines_ci$del_qt <= max_del & lines_ci$phi_qt <= max_phi, ]

df$m <- 0
df$m[df$firmtype == "CD"] <- df$prob.coastal.domestic[df$firmtype == "CD"]
df$m[df$firmtype == "CI"] <- df$prob.coastal.exporter[df$firmtype == "CI"]
df$m[df$firmtype == "PD"] <- df$prob.inland.domestic[df$firmtype == "PD"]
df$m[df$firmtype == "PI"] <- df$prob.inland.exporter[df$firmtype == "PI"]


var <- "dirtyinput"
df$dirtyinput <- df$d.inland + df$d.coastal

scale_range <- c(1.5, 20)
df$m[df$type == "Exit"] <- 0

# CREATE FILL VARIABLES (TOP AND BOTTOM 10%)
df$fill_var <- 0
vals <- quantile(df[df$type != "Exit", var], probs = seq(0, 1, 0.1))[c(2, 10)]
pval <- order(df[df$firmtype %in% c("PD", "PI"), var], decreasing = TRUE)[20]
pval <- df[df$firmtype %in% c("PD", "PI"), var][pval]
cval <- order(df[df$firmtype %in% c("CD", "CI"), var], decreasing = TRUE)[20]
cval <- df[df$firmtype %in% c("CD", "CI"), var][cval]
df$fill_var[df$type != "Exit" & df[, var] > vals[2]] <- 1 
df$fill_var[df$type != "Exit" & df[, var] < vals[1]] <- 1


# SHADE IN EXITING FIRMS AREA
m <- df[df$type == "Exit", c('phi_qt', 'del_qt')]
m$phi_qt <- m$phi_qt * scale
m$del_qt <- m$del_qt * scale
min_del <- min(m$del_qt)
max_del <- max(m$del_qt)
m$del_qt[m$del_qt == min_del] <- round(m$del_qt[m$del_qt == min_del]) - scale/2
m$del_qt[m$del_qt == max_del] <- floor(m$del_qt[m$del_qt == max_del]) + scale/2

hull <- chull(m) #create coordinates of convex hull
coords <- m[ c( hull, hull[1]), ]

# LABEL THE EXITING FIRMS
exit_x <- mean(df[df$type == "Exit", "del_qt"]) * scale
exit_y <- mean(df[df$type == "Exit", "phi_qt"]) * scale

# SHADE IN THE TOP 20% OF PERIPHERY FIRMS 
m_p <- df[df$firmtype %in% c("PD", "PI") & df[, var] >= pval, c('phi_qt', 'del_qt')]
m_p$phi_qt <- m_p$phi_qt * scale
m_p$del_qt <- m_p$del_qt * scale
hull <- chull(m_p) #create coordinates of convex hull
coords_p <- m_p[c( hull, hull[1]), ]
m_p$id <- "20 Largest\nPeriphery Firms"
m_p$x0 <- 30
m_p$y0 <- 80

# LABEL TOP 10% FIRMS AREA
lb_py <- mean(m_p$phi_qt)
lb_px <- mean(m_p$del_qt)

# SHADE IN THE TOP 10% OF CORE FIRMS 
m_c <- df[df$firmtype %in% c("CD", "CI") & df[, var] >= cval, c('phi_qt', 'del_qt')]
m_c$phi_qt <- m_c$phi_qt * scale
m_c$del_qt <- m_c$del_qt * scale
hull <- chull(m_c) #create coordinates of convex hull
coords_c <- m_c[c(hull, hull[1]), ]
m_c$id <- "20 Largest\nCore Firms"
m_c$x0 <- 85
m_c$y0 <- 60

# LABEL TOP 10% FIRMS AREA
lb_cy <- mean(m_c$phi_qt)
lb_cx <- mean(m_c$del_qt)

m <- rbind(m_p, m_c)

df <- merge(df, m, by = c("phi_qt", "del_qt"), all.x = T)

label_parse <- function(breaks) {
   parse(text = breaks)
}

firm_plot <- function(var, scale_range, out_path, trans, q_lab, pfix) {
    # CREATE LEGEND ORDER
    ordered_leg <- c("Exit", "Periphery-Domestic",
        "Periphery-International", "Core-Domestic", "Core-International")

    color_leg <- c("Exit" = 'grey', "Periphery-Domestic" = "#9370DB",
                    "Periphery-International" = "#0F2080", "Core-Domestic" = "#7ac77f", #"#A95AA1" ,
                    "Core-International" = "#cc3137")
    shape_leg <- c("Exit" = 21, "Periphery-Domestic" = 23, "Periphery-International" = 8,
                    "Core-Domestic" = 22, "Core-International" = 24)


    n_obs <- length(df[df$type != "Exit" & df[,var] != 0 & !is.na(df[,var]), var])

    mylabels <- function(breaks){
        plot_var <- quantile(df[df$type != "Exit" & df[,var] != 0, var], seq(0, 1, length.out = n_obs), na.rm = T)
        qs_labs <- cut(breaks, plot_var, include.lowest=TRUE, label = names(plot_var[2:length(plot_var)]))
        qs_labs <- as.numeric(sub("%", "", qs_labs))
        qs_labs[is.na(qs_labs)] <- 0
        labels <- sprintf("%.0f%%", qs_labs) # make your labels here
        return(labels)
    }
    
    
    p <- ggplot(data = df[df$type != "Exit", ], # DO NOT PLOT EXITING FIRMS
            aes(del_qt * scale, phi_qt * scale, fill = factor(type),
                size = .data[[var]], color = type)) + #shape = type, #.data[[var]]
        geom_point(stroke = 1.25) +
        geom_path(data = lines_cd, aes(del_qt*scale, phi_qt*scale), linetype = "dashed",
                inherit.aes = F, linewidth = 1.2, color = "#b00f14") +
        geom_path(data = lines_ci, aes(del_qt*scale, phi_qt*scale), linetype = "dashed",
                inherit.aes = F, linewidth = 1.2, color = "forestgreen") +
        geom_mark_hull(data = m,
            aes(x0 = x0, y0 = y0, x = del_qt, y = phi_qt, 
                label = id, fill = id),
            alpha = .25, inherit.aes = F, concavity = 2,
            expand = unit(2.5, "mm"),
            con.border = "all",
            linetype = 0,
            label.fontsize = 14,
            label.family = "Palatino",
            con.arrow = arrow(length = unit(3, "mm")),
            label.margin = margin(1, 1, 1, 1, "mm"),
            con.size = 0.75, con.cap = 0) +
        scale_color_manual(values = color_leg,
            breaks = ordered_leg,
            guide = guide_legend(override.aes = list(size = 4))) +
        scale_fill_manual(values = color_leg,
            breaks = ordered_leg) +
        geom_polygon(data = coords, aes(del_qt, phi_qt), color = "darkgrey",
                    size = 1, linetype = "dashed", #fill = NULL, 
                    alpha = 0.0, inherit.aes = F) +
        theme_bw() +
        guides(fill = "none", alpha ='none', size = "none") + #
        annotate("text", x = exit_x, y = exit_y,
            label = "Exit", size = 5, family = "Palatino") + # LABEL EXITING FIRMS
        annotate("label", x = 45, y = 12, color = "#9370DB", 
            label = "Periphery\nDomestic", size = 5, family = "Palatino") + # LABEL Periphery Domestic FIRMS
        annotate("label", x = 83, y = 25, color = 'forestgreen',
            label = "Core Domestic", size = 5, family = "Palatino") + # LABEL Periphery Domestic FIRMS
        annotate("label", x = 58, y = 100, color = "#b00f14", 
            label = "Core International", size = 5, family = "Palatino") +
        ylab(expression("Efficiency" ~ italic(phi))) +  #\U03C6 '"Efficiency" ~ phi' \U0001D719
        xlab(expression("Size" ~ italic(delta))) +  #\U03B4 expression(Size- ~"\U0001D6FF"))
        labs(size = q_lab, shape = NULL, color = NULL,) +
        theme(legend.position="bottom",
            legend.box = "vertical",
            legend.box.margin = margin(3, 3, 3, 3),
            text = element_text(size = 16, family = "Palatino"), #
            axis.title = element_text(size = 16), 
            axis.text = element_text(size = 16, color = "black"),
            legend.title = element_text(size = 16)) #+ 


        showtext_auto()
        if (var == "m" | var == "m_alt") {
            mylegend <- g_legend(p)
            tmp1 <- file.path(out_path, "location_legend.eps")
            
            cairo_ps(tmp1)
            grid.draw(mylegend)
            dev.off()

            qs <- quantile(df[df$type != "Exit" & df[,var] != 0, var])
            names(qs) <- paste0(round(qs, 2) * 100, "%")
            p <- p + scale_size(labels = names(qs),#qs_labs,
                    breaks = qs,
                    range = scale_range,
                    trans = trans,
                    guide = guide_legend(override.aes = list(
                            shape = 16, 
                            color = "black", fill = "black"))) + 
                guides(size = guide_legend(nrow = 1), color = "none")
        } else {
            p <- p + scale_size(labels = mylabels,
                    breaks = quantile(df[df$type != "Exit" & df[,var] != 0 & !is.na(df[,var]), var])[2:5],
                    range = scale_range,
                    trans = trans,
                    guide = guide_legend(override.aes = list(
                            shape = 16,
                            color = "black", fill = "black"))) + 
                guides(size = guide_legend(nrow = 1), color = "none")
        }

    # SAVE FIGURES
    fn <- paste0(pfix, var, "_qt_FINAL.eps")
    
    cairo_ps(file.path(out_path, fn), width = 8, height = 8.25, fallback_resolution = 320)
    print(p)
    dev.off()
    # ggsave(file.path(out_path, fn), 
    #         plot = p, device="cairo_ps",
    #         dpi = 320, width = 8, height = 8.25, units = "in")
    showtext_auto(FALSE)
}


PREFIXES <- list("emission" = "fig06a_",
                "intensity" = "fig06b_",
                "dirtyinput" = "fig07a_",
                "m" = "fig05b_",
                "fa" = "fig07b_",
                "dirty_share" = "fig08b_",
                'abate_shr_dirty_input_price' = "fig08a_")

VARIABLES <- names(PREFIXES)

scales <- list("emission" = c(2.25, 10),
                "intensity" = c(0.25, 10),
                "dirtyinput" = c(2.25, 10),
                "abate_ttlexp_share" = c(0.25, 10), 
                "dxratio" = c(0.25, 10),
                "m" = c(0.25, 10),
                "fa" = c(2.25, 10),
                "dirty_share" = c(2.25,10),
                "abate_shr_dirty_input_price" = c(2.25,10), 
                "dirty_exp" = c(2.25, 10)) 

trans_list <- list("emission" = "identity",
                    "intensity" = "identity",
                    "dirtyinput" = "identity",
                    "abate_ttlexp_share" = "identity",
                    "dxratio" = "identity",
                    'm' = 'identity',
                    'fa' = 'identity',
                    "dirty_share" = "identity",
                    "dirty_exp" = "identity",
                    'abate_shr_dirty_input_price' = "identity")



df$emission <- (df$opt.e.coastal + df$opt.e.inland)
df$emission[df$firm_loc == "prob.inland"] <- df$opt.e.inland[df$firm_loc == "prob.inland"]
df$emission[df$firm_loc == "prob.coastal"] <- df$opt.e.coastal[df$firm_loc == "prob.coastal"]


df$intensity <- (df$coastal.eint + df$inland.eint)
df$intensity[df$firm_loc == "prob.inland"] <- df$inland.eint[df$firm_loc == "prob.inland"]
df$intensity[df$firm_loc == "prob.coastal"] <- df$coastal.eint[df$firm_loc == "prob.coastal"]


df$abate_ttlexp_share <- (df$share.fa.coastal + df$share.fa.inland)
df$abate_ttlexp_share[df$firm_loc == "prob.inland"] <- df$share.fa.inland[df$firm_loc == "prob.inland"]
df$abate_ttlexp_share[df$firm_loc == "prob.coastal"] <- df$share.fa.coastal[df$firm_loc == "prob.coastal"]
abate_shr_lb <- min(df$abate_ttlexp_share)
df$abate_ttlexp_share[df$fa.inland * df$prob.inland + df$fa.coastal * df$prob.coastal < 0.0001] <- 0


df$fa <- (df$fa.coastal + df$fa.inland)
df$fa[df$firm_loc == "prob.inland"] <- df$fa.inland[df$firm_loc == "prob.inland"]
df$fa[df$firm_loc == "prob.coastal"] <- df$fa.coastal[df$firm_loc == "prob.coastal"]


df$Kappa <- (df$Kappa.coastal + df$Kappa.inland)
df$Kappa[df$firm_loc == "prob.inland"] <- df$Kappa.inland[df$firm_loc == "prob.inland"]
df$Kappa[df$firm_loc == "prob.coastal"] <- df$Kappa.coastal[df$firm_loc == "prob.coastal"]



df$X <- (df$coastal.X + df$inland.X)
df$X[df$firm_loc == "prob.inland"] <- df$inland.X[df$firm_loc == "prob.inland"]
df$X[df$firm_loc == "prob.coastal"] <- df$coastal.X[df$firm_loc == "prob.coastal"]


df$dxratio <- (df$dxratio.inland + df$dxratio.coastal)
df$dxratio[df$firm_loc == "prob.inland"] <- df$dxratio.inland[df$firm_loc == "prob.inland"]
df$dxratio[df$firm_loc == "prob.coastal"] <- df$dxratio.coastal[df$firm_loc == "prob.coastal"]


df$dirty_share <- (df$share.d.inland + df$share.d.coastal)
df$dirty_share[df$firm_loc == "prob.inland"] <- df$share.d.inland[df$firm_loc == "prob.inland"]
df$dirty_share[df$firm_loc == "prob.coastal"] <- df$share.d.coastal[df$firm_loc == "prob.coastal"]

df$dirtyinput <- df$dirtyinput
df$dirtyinput[df$firm_loc == "prob.inland"] <- df$d.inland[df$firm_loc == "prob.inland"]
df$dirtyinput[df$firm_loc == "prob.coastal"] <- df$d.coastal[df$firm_loc == "prob.coastal"]

df$int.x <- df$int.x.coastal + df$int.x.inland
df$int.x[df$firm_loc == "prob.inland"] <- df$int.x.inland[df$firm_loc == "prob.inland"]
df$int.x[df$firm_loc == "prob.coastal"] <- df$int.x.coastal[df$firm_loc == "prob.coastal"]


df$Phi <- df$Phi.coastal + df$Phi.inland
df$Phi[df$firm_loc == "prob.inland"] <- df$Phi.inland[df$firm_loc == "prob.inland"]
df$Phi[df$firm_loc == "prob.coastal"] <- df$Phi.coastal[df$firm_loc == "prob.coastal"]


df$d <- df$d.coastal + df$d.inland
df$d[df$firm_loc == "prob.inland"] <- df$d[df$firm_loc == "prob.inland"]
df$d[df$firm_loc == "prob.coastal"] <- df$d[df$firm_loc == "prob.coastal"]


df$X <- df$coastal.X + df$inland.X
df$X[df$firm_loc == "prob.inland"] <- df$inland.X[df$firm_loc == "prob.inland"]
df$X[df$firm_loc == "prob.coastal"] <- df$coastal.X[df$firm_loc == "prob.coastal"]

df$d <- df$d.coastal + df$d.inland
df$d[df$firm_loc == "prob.inland"] <- df$d.inland[df$firm_loc == "prob.inland"]
df$d[df$firm_loc == "prob.coastal"] <- df$d.coastal[df$firm_loc == "prob.coastal"]


df$dirty_exp <- (df$Kappa) / (df$X * df$Phi^(-1/3) + df$fa)
df$dirty_exp <- (df$Kappa) / (df$Phi^(-1/3) + df$fa/df$X)
df$abate_ttlexp_share <- (df$fa)/df$d / (df$X * df$Phi^(-1/3) + df$fa)
df$abate_ttlexp_share <- (df$fa)/df$d / (df$Phi^(-1/3) + df$fa/df$X)
df$abate_ttlexp_share <- ((df$phi / (df$fa/df$d)) ^ (3))/df$Phi


# What would the unit cost share be if abatement were set to the min val
min_fa <- min(df$fa[df$fa > 0], na.rm = T)
df$fa_cf <- min_fa/100
df$Kappa_cf <- w + (t * epsilon) / min_fa
df$Kappa_cf[df$firm_loc == "prob.coastal"] <- (df$Kappa_cf * df$prob.coastal)[df$firm_loc == "prob.coastal"]
df$Kappa_cf[df$firm_loc == "prob.inland"] <- (df$Kappa_cf * df$prob.inland)[df$firm_loc == "prob.inland"]
df$Kappa_cf[df$firm_loc == "Exit"] <- 0

df$abated_Kappa <- df$Kappa_cf-df$Kappa
df$abate_shr_dirty_input_price <- ((df$phi / df$abated_Kappa) ^ (3))/df$Phi
df$dirty_share <- ((df$phi / df$Kappa) ^ (3))/df$Phi
df$abate_shr_dirty_input_price <- df$fa/df$X * (df$Phi)^(-1/3)


for (variable in VARIABLES) {
    print(variable)
    scale_range <- scales[[variable]]
    trans <- trans_list[[variable]]
    pfix <- PREFIXES[[variable]]
    if (variable == "m") {
        q_lab <- "Max Location Prob."
    } else {
        q_lab <- "Percentile"
    }
    firm_plot(variable, scale_range, out_path, trans, q_lab, pfix)
}

####################################################################
###################### FIGURE 4(a) #################################
####################################################################

objects <- all_results[, , 9]

phi <- objects$phi
del <- objects$del
n_grid <- length(phi)


############## GET GCHI DATA!!!! ##################
M = c(mat$m.phi, mat$m.del)                                        
V = matrix(c(mat$v.phi^2, mat$v.cor*mat$v.phi*mat$v.del, mat$v.cor*mat$v.phi*mat$v.del, mat$v.del^2), ncol = 2) 
         
g_chi_fun <- function(phi,del) dmvnorm(c(log(phi), log(del)), M, sigma =V, log = FALSE, checkSymmetry = TRUE)/(phi*del)
g_chi <- matrix(rep(0, n_grid*n_grid), ncol = n_grid)
for (n_phii in seq_len(n_grid)) {
    for (n_dell in seq_len(n_grid)) {
        g_chi[n_phii, n_dell] <- g_chi_fun(phi[n_phii], del[n_dell])
    }
}
g_chi <- g_chi/sum(g_chi)


# initialize our data frame
df <- data.frame("phi" = rep(phi, n_grid),
                "del" = rep(del, each = n_grid),
                "g_chi" = c(g_chi))

vars <- c("prob.coastal", "prob.inland",
    "prob.inland.domestic", "prob.inland.exporter",
    "prob.coastal.domestic", "prob.coastal.exporter",
    "opt.e.coastal", "opt.e.inland",
    "inland.X", "coastal.X", "d.inland", "d.coastal",
    "share.fa.inland", "share.fa.coastal",
    "dxratio.inland", "dxratio.coastal",
    "xdratio.inland", "xdratio.coastal")

for (var in vars) {
    df[var] <- c(objects[[var]])
}
df$loc_prob_max <- apply(df[, loc_vars], 1, max, na.rm=TRUE)
df$firm_loc <- colnames(df[, loc_vars])[apply(df[, loc_vars], 1, which.max)]

df$firm_loc <- "prob.coastal"
df$firm_loc[df$prob.inland >= il_prob] <- "prob.inland"

df$firm_loc[df$loc_prob_max == 0] <- "Exit"

intl_vars <- c("prob.inland.domestic", "prob.inland.exporter")
df$intl_inland <- colnames(df[, intl_vars])[apply(df[, intl_vars], 1, which.max)]
intl_vars <- c("prob.coastal.domestic", "prob.coastal.exporter")
df$intl_coastal <- colnames(df[, intl_vars])[apply(df[, intl_vars], 1, which.max)]

df$firmtype <- "Exit"
df$firmtype[df$firm_loc == "prob.inland"] <- df$intl_inland[df$firm_loc == "prob.inland"]
df$firmtype[df$firm_loc == "prob.coastal"] <- df$intl_coastal[df$firm_loc == "prob.coastal"]

df$firmtype[df$firmtype == "prob.inland.domestic"] <- "PD"
df$firmtype[df$firmtype == "prob.inland.exporter"] <- "PI"
df$firmtype[df$firmtype == "prob.coastal.domestic"] <- "CD"
df$firmtype[df$firmtype == "prob.coastal.exporter"] <- "CI"

df$m <- df$max

df$type = "Exit"
df$type[df$firmtype == "PD"] <- "Periphery-Domestic"
df$type[df$firmtype == "PI"] <- "Periphery-International"
df$type[df$firmtype == "CI"] <- "Core-International"
df$type[df$firmtype == "CD"] <- "Core-Domestic"


# get quantiles of phi and del
phis <- unique(df$phi)
phi_qt <- data.frame(phi = phis, phi_qt = seq_along(phis))
dels <- unique(df$del)
del_qt <- data.frame(del = dels, del_qt = seq_along(dels))
df <- merge(df, phi_qt, by = "phi")
df <- merge(df, del_qt, by = "del")

scale <- 100 / max(df$phi_qt, df$del_qt)

phi_pts <- c()
for (del in unique(df$del_qt)) {
    v1 <- min(df[df$del_qt == del & df$firmtype == "CI", "phi_qt"])
    v2 <- max(df[df$del_qt == del & df$firmtype == "CD", "phi_qt"])
    phi_pts <- c(phi_pts, mean(v1, v2))
}

del_pts_grid <- unique(df$del_qt)

# draw line between cd and ci 
del_pts <- c()
for (phi in unique(df$phi_qt)) {
    v1 <- min(df[df$phi_qt == phi & df$firmtype == "CI", "del_qt"])
    v2 <- max(df[df$phi_qt == phi & df$firmtype == "CD", "del_qt"])
    del_pts <- c(del_pts, mean(v1, v2))
}

phi_pts_grid <- unique(df$phi_qt)
lines_1 <- data.frame(del_qt = del_pts, phi_qt = phi_pts_grid)
lines_2 <- data.frame(del_qt = del_pts_grid, phi_qt = phi_pts)
lines_cd <- unique(union(lines_1, lines_2))
lines_cd <- lines_cd[order(lines_cd$del_qt, -lines_cd$phi_qt), ]
lines_cd$del_qt <- lines_cd$del_qt - 0.5
lines_cd$phi_qt <- lines_cd$phi_qt - 0.5

# draw line between pd and cd
phi_pts <- c()
for (del in unique(df$del_qt)) {
    v1 <- min(df[df$del_qt == del & df$firmtype == "CD", "phi_qt"])
    v2 <- max(df[df$del_qt == del & df$firmtype == "PD", "phi_qt"])
    phi_pts <- c(phi_pts, mean(v1, v2))
}

del_pts_grid <- unique(df$del_qt)

# draw line between pd and cd 
del_pts <- c()
for (phi in unique(df$phi_qt)) {
    v1 <- min(df[df$phi_qt == phi & df$firmtype == "CD", "del_qt"])
    v2 <- max(df[df$phi_qt == phi & df$firmtype == "PD", "del_qt"])
    del_pts <- c(del_pts, mean(v1, v2))
}

phi_pts_grid <- unique(df$phi_qt)
lines_1 <- data.frame(del_qt = del_pts, phi_qt = phi_pts_grid)
lines_2 <- data.frame(del_qt = del_pts_grid, phi_qt = phi_pts)
lines_ci <- unique(union(lines_1, lines_2))
lines_ci <- lines_ci[order(lines_ci$del_qt, -lines_ci$phi_qt), ]
lines_ci$del_qt <- lines_ci$del_qt - 0.5
lines_ci$phi_qt <- lines_ci$phi_qt - 0.5

max_phi <- max(df[df$firmtype == "PD", "phi_qt"])+1
max_del <- max(df[df$firmtype == "PD", "del_qt"])+1

lines_ci <- lines_ci[lines_ci$del_qt <= max_del & lines_ci$phi_qt <= max_phi, ]

df$m <- 0
df$m[df$firmtype == "CD"] <- df$prob.coastal.domestic[df$firmtype == "CD"]
df$m[df$firmtype == "CI"] <- df$prob.coastal.exporter[df$firmtype == "CI"]
df$m[df$firmtype == "PD"] <- df$prob.inland.domestic[df$firmtype == "PD"]
df$m[df$firmtype == "PI"] <- df$prob.inland.exporter[df$firmtype == "PI"]



var <- "dirtyinput"
df$dirtyinput <- df$d.inland + df$d.coastal

scale_range <- c(1.5, 20)
df$m[df$type == "Exit"] <- 0

# CREATE FILL VARIABLES (TOP AND BOTTOM 10%)
df$fill_var <- 0
vals <- quantile(df[df$type != "Exit", var], probs = seq(0, 1, 0.1))[c(2, 10)]
pval <- order(df[df$firmtype %in% c("PD", "PI"), var], decreasing = TRUE)[20]
pval <- df[df$firmtype %in% c("PD", "PI"), var][pval]
cval <- order(df[df$firmtype %in% c("CD", "CI"), var], decreasing = TRUE)[20]
cval <- df[df$firmtype %in% c("CD", "CI"), var][cval]
df$fill_var[df$type != "Exit" & df[, var] > vals[2]] <- 1 
df$fill_var[df$type != "Exit" & df[, var] < vals[1]] <- 1


# SHADE IN EXITING FIRMS AREA
m <- df[df$type == "Exit", c('phi_qt', 'del_qt')]
m$phi_qt <- m$phi_qt * scale
m$del_qt <- m$del_qt * scale
min_del <- min(m$del_qt)
max_del <- max(m$del_qt)
m$del_qt[m$del_qt == min_del] <- round(m$del_qt[m$del_qt == min_del]) - scale/2
m$del_qt[m$del_qt == max_del] <- floor(m$del_qt[m$del_qt == max_del]) + scale/2

hull <- chull(m) #create coordinates of convex hull
coords <- m[ c( hull, hull[1]), ]

# LABEL THE EXITING FIRMS
exit_x <- mean(df[df$type == "Exit", "del_qt"]) * scale
exit_y <- mean(df[df$type == "Exit", "phi_qt"]) * scale


# SHADE IN THE TOP 20% OF PERIPHERY FIRMS 
m_p <- df[df$firmtype %in% c("PD", "PI") & df[, var] >= pval, c('phi_qt', 'del_qt')]
m_p$phi_qt <- m_p$phi_qt * scale
m_p$del_qt <- m_p$del_qt * scale
hull <- chull(m_p) #create coordinates of convex hull
coords_p <- m_p[c( hull, hull[1]), ]
m_p$id <- "20 Largest\nPeriphery Firms"
m_p$x0 <- 30
m_p$y0 <- 80

# LABEL TOP 10% FIRMS AREA
lb_py <- mean(m_p$phi_qt)
lb_px <- mean(m_p$del_qt)

# SHADE IN THE TOP 10% OF CORE FIRMS 
m_c <- df[df$firmtype %in% c("CD", "CI") & df[, var] >= cval, c('phi_qt', 'del_qt')]
m_c$phi_qt <- m_c$phi_qt * scale
m_c$del_qt <- m_c$del_qt * scale
hull <- chull(m_c) #create coordinates of convex hull
coords_c <- m_c[c(hull, hull[1]), ]
m_c$id <- "20 Largest\nCore Firms"
m_c$x0 <- 85
m_c$y0 <- 60

# LABEL TOP 10% FIRMS AREA
lb_cy <- mean(m_c$phi_qt)
lb_cx <- mean(m_c$del_qt)

m <- rbind(m_p, m_c)

df <- merge(df, m, by = c("phi_qt", "del_qt"), all.x = T)

df$m_alt <- df$m

firm_plot("m_alt", c(1, 9.5), out_path, "identity", "Max Location Prob.", "fig05a_")